from tkinter import filedialog

import cv2
import numpy as np
from pathlib import Path

# 3. 定义红色和蓝色的HSV范围
lower_red1 = np.array([0, 70, 50])
upper_red1 = np.array([10, 255, 255])
lower_red2 = np.array([170, 70, 50])
upper_red2 = np.array([180, 255, 255])

lower_blue = np.array([100, 150, 0])
upper_blue = np.array([140, 255, 255])


kernel = np.ones((3, 3), np.uint8)

directory = filedialog.askdirectory()
folder_path = Path(directory)
images = folder_path.rglob("*.jpg")


for image in images:

    # 2. 转换颜色空间到HSV
    image_name = directory + "/" + image.name
    img = cv2.imread(image_name)
    hsv_image = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    # 4. 创建掩码，识别红色和蓝色
    red_mask1 = cv2.inRange(hsv_image, lower_red1, upper_red1)
    red_mask2 = cv2.inRange(hsv_image, lower_red2, upper_red2)
    blue_mask = cv2.inRange(hsv_image, lower_blue, upper_blue)

    # 5. 合并红色掩码（如果你认为红色可能跨越两个色调）
    red_mask = cv2.bitwise_or(red_mask1, red_mask2)

    # 6. 应用掩码到原始图像，提取红色和蓝色区域
    red_image = cv2.bitwise_and(img, img, mask=red_mask)
    blue_image = cv2.bitwise_and(img, img, mask=blue_mask)

    # 7. 可选：形态学操作去除噪声
    red_image = cv2.morphologyEx(red_image, cv2.MORPH_OPEN, kernel)
    blue_image = cv2.morphologyEx(blue_image, cv2.MORPH_OPEN, kernel)

    red_draw = cv2.morphologyEx(red_mask, cv2.MORPH_OPEN, kernel)

    # 查找轮廓
    contours, _ = cv2.findContours(red_draw, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 8. 报出有红色绘图的图片名字
    if len(contours) > 0:
        print("have,but not sure." + str(image.name))
        for contour in contours:
            box = cv2.boundingRect(contour)
            if box[2] * box[3] > 50:
                print(box[2])
                print(box[3])
                print(image)
